#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
流式TTS服务模块

基于火山引擎的WebSocket流式语音合成服务
支持实时文本转语音，边输入边合成
"""

import asyncio
import base64
import gzip
import json
import uuid
import websockets
from typing import Optional, Callable, Dict, Any
import copy
import time


# 协议常量定义
MESSAGE_TYPES = {
    0x01: "full client request",
    0x02: "audio-only client request", 
    0x09: "full server response",
    0x0b: "audio-only server response", 
    0x0c: "frontend server response",
    0x0f: "error message from server"
}

MESSAGE_TYPE_SPECIFIC_FLAGS = {
    0: "no sequence number",
    1: "sequence number > 0",
    2: "last message from server (seq < 0)",
    3: "sequence number < 0"
}

MESSAGE_SERIALIZATION_METHODS = {
    0: "no serialization",
    1: "JSON",
    15: "custom type"
}

MESSAGE_COMPRESSIONS = {
    0: "no compression",
    1: "gzip",
    15: "custom compression method"
}


class StreamingTTSService:
    """
    流式TTS服务
    
    实现实时文本转语音，边输入边合成
    支持异步处理和回调机制
    """
    
    def __init__(self,
                 appid: str = "1871057448",
                 token: str = "tnph6TBEuRv4MWgh425-XUf6K7AcS2_5",
                 cluster: str = "volcano_tts",
                 host: str = "openspeech.bytedance.com",
                 **kwargs):
        """
        初始化流式TTS服务
        
        Args:
            appid: 火山引擎应用ID
            token: 访问令牌
            cluster: 集群名称
            host: 服务主机地址
            **kwargs: 其他配置参数
        """
        self.appid = appid
        self.token = token
        self.cluster = cluster
        self.host = host
        self.api_url = f"wss://{host}/api/v1/tts/ws_binary"
        
        # 音频参数
        self.voice_type = kwargs.get("voice_type", "BV009_streaming")  # 流式语音类型
        self.encoding = kwargs.get("encoding", "mp3")  # 音频编码格式
        self.speed_ratio = kwargs.get("speed_ratio", 1.0)  # 语速
        self.volume_ratio = kwargs.get("volume_ratio", 1.0)  # 音量
        self.pitch_ratio = kwargs.get("pitch_ratio", 1.0)  # 音调
        
        # 用户标识
        self.uid = kwargs.get("uid", "streaming_tts_user")
        
        # 协议头 - version: b0001, header size: b0001, message type: b0001 (Full client request)
        # message type specific flags: b0000, serialization: b0001 (JSON), compression: b0001 (gzip)
        self.default_header = bytearray(b'\x11\x10\x11\x00')
        
        # WebSocket连接
        self.ws = None
        self.is_running = False
        
        # 回调函数
        self.on_audio_data: Optional[Callable] = kwargs.get('on_audio_data', None)  # 音频数据回调
        self.on_complete: Optional[Callable] = kwargs.get('on_complete', None)  # 完成回调
        self.on_error: Optional[Callable] = kwargs.get('on_error', None)  # 错误回调
        
        # 统计信息
        self.audio_chunks_count = 0
        self.total_audio_size = 0
        
        print(f"🔊 流式TTS服务初始化完成")
        print(f"   AppID: {appid}")
        print(f"   集群: {cluster}")
        print(f"   语音类型: {self.voice_type}")
        print(f"   编码格式: {self.encoding}")
    
    def construct_request(self, text: str, operation: str = "submit") -> dict:
        """
        构建TTS请求
        
        Args:
            text: 要合成的文本
            operation: 操作类型 (submit: 提交合成, query: 查询结果)
        
        Returns:
            请求字典
        """
        reqid = str(uuid.uuid4())
        request = {
            "app": {
                "appid": self.appid,
                "token": "access_token",  # 占位符，实际通过header传递
                "cluster": self.cluster
            },
            "user": {
                "uid": self.uid
            },
            "audio": {
                "voice_type": self.voice_type,
                "encoding": self.encoding,
                "speed_ratio": self.speed_ratio,
                "volume_ratio": self.volume_ratio,
                "pitch_ratio": self.pitch_ratio,
            },
            "request": {
                "reqid": reqid,
                "text": text,
                "text_type": "plain",
                "operation": operation
            }
        }
        return request
    
    def build_request_payload(self, text: str, operation: str = "submit") -> bytearray:
        """
        构建请求payload
        
        Args:
            text: 要合成的文本
            operation: 操作类型
        
        Returns:
            完整的请求字节数组
        """
        # 构建请求JSON
        request_json = self.construct_request(text, operation)
        
        # 序列化并压缩
        payload_bytes = str.encode(json.dumps(request_json))
        payload_bytes = gzip.compress(payload_bytes)
        
        # 构建完整请求
        full_request = bytearray(self.default_header)
        full_request.extend((len(payload_bytes)).to_bytes(4, 'big'))  # payload size
        full_request.extend(payload_bytes)  # payload
        
        return full_request
    
    def parse_response(self, res: bytes) -> Dict[str, Any]:
        """
        解析服务器响应
        
        Args:
            res: 响应字节数据
        
        Returns:
            解析后的响应字典
        """
        # 解析协议头
        protocol_version = res[0] >> 4
        header_size = res[0] & 0x0f
        message_type = res[1] >> 4
        message_type_specific_flags = res[1] & 0x0f
        serialization_method = res[2] >> 4
        message_compression = res[2] & 0x0f
        reserved = res[3]
        header_extensions = res[4:header_size*4]
        payload = res[header_size*4:]
        
        result = {
            'protocol_version': protocol_version,
            'header_size': header_size,
            'message_type': message_type,
            'message_type_specific_flags': message_type_specific_flags,
            'serialization_method': serialization_method,
            'message_compression': message_compression,
            'audio_data': None,
            'error': None,
            'is_last': False
        }
        
        # 处理不同类型的消息
        if message_type == 0x0b:  # audio-only server response
            if message_type_specific_flags == 0:  # no sequence number (ACK)
                result['is_ack'] = True
                return result
            else:
                # 有音频数据
                sequence_number = int.from_bytes(payload[:4], "big", signed=True)
                payload_size = int.from_bytes(payload[4:8], "big", signed=False)
                audio_data = payload[8:]
                
                result['sequence_number'] = sequence_number
                result['audio_data'] = audio_data
                result['audio_size'] = payload_size
                result['is_last'] = (sequence_number < 0)  # 负数表示最后一帧
                
        elif message_type == 0x0f:  # error message
            code = int.from_bytes(payload[:4], "big", signed=False)
            msg_size = int.from_bytes(payload[4:8], "big", signed=False)
            error_msg = payload[8:]
            
            if message_compression == 1:
                error_msg = gzip.decompress(error_msg)
            error_msg = str(error_msg, "utf-8")
            
            result['error'] = {
                'code': code,
                'message': error_msg
            }
            result['is_last'] = True
            
        elif message_type == 0x0c:  # frontend server response
            msg_size = int.from_bytes(payload[:4], "big", signed=False)
            payload = payload[4:]
            if message_compression == 1:
                payload = gzip.decompress(payload)
            result['frontend_message'] = payload
        
        return result
    
    async def connect(self):
        """建立WebSocket连接"""
        try:
            header = {"Authorization": f"Bearer; {self.token}"}
            self.ws = await websockets.connect(
                self.api_url,
                additional_headers=header,
                ping_interval=None,
                max_size=100000000  # 100MB
            )
            self.is_running = True
            print("✅ TTS WebSocket连接建立成功")
            return True
            
        except Exception as e:
            print(f"❌ 连接异常: {e}")
            if self.on_error:
                self.on_error(str(e))
            return False
    
    async def synthesize_text(self, text: str) -> bool:
        """
        合成文本为语音
        
        Args:
            text: 要合成的文本
        
        Returns:
            是否成功
        """
        if not self.ws or not self.is_running:
            print("❌ WebSocket未连接")
            return False
        
        try:
            print(f"\n📝 开始合成文本: {text[:50]}...")
            
            # 重置统计
            self.audio_chunks_count = 0
            self.total_audio_size = 0
            
            # 构建并发送请求
            request_payload = self.build_request_payload(text, "submit")
            await self.ws.send(request_payload)
            print("📤 请求已发送")
            
            # 接收响应
            while self.is_running:
                try:
                    res = await self.ws.recv()
                    result = self.parse_response(res)
                    
                    # 处理音频数据
                    if result.get('audio_data'):
                        self.audio_chunks_count += 1
                        self.total_audio_size += len(result['audio_data'])
                        
                        # 触发音频数据回调
                        if self.on_audio_data:
                            await self.on_audio_data({
                                'audio_data': result['audio_data'],
                                'sequence_number': result.get('sequence_number'),
                                'is_last': result['is_last'],
                                'chunk_index': self.audio_chunks_count,
                                'total_size': self.total_audio_size
                            })
                        
                        # 如果是最后一帧，结束
                        if result['is_last']:
                            print(f"✅ 合成完成 - 共 {self.audio_chunks_count} 个音频块, 总大小 {self.total_audio_size} 字节")
                            if self.on_complete:
                                await self.on_complete({
                                    'total_chunks': self.audio_chunks_count,
                                    'total_size': self.total_audio_size
                                })
                            return True
                    
                    # 处理错误
                    if result.get('error'):
                        error = result['error']
                        error_msg = f"TTS错误 (code={error['code']}): {error['message']}"
                        print(f"❌ {error_msg}")
                        if self.on_error:
                            self.on_error(error_msg)
                        return False
                    
                    # 处理ACK
                    if result.get('is_ack'):
                        print("📬 收到ACK")
                        continue
                    
                except Exception as e:
                    print(f"❌ 接收响应异常: {e}")
                    if self.on_error:
                        self.on_error(str(e))
                    return False
            
            return True
            
        except Exception as e:
            print(f"❌ 合成异常: {e}")
            if self.on_error:
                self.on_error(str(e))
            return False
    
    async def disconnect(self):
        """断开WebSocket连接"""
        self.is_running = False
        if self.ws:
            await self.ws.close()
            self.ws = None
            print("🔌 TTS WebSocket连接已关闭")


class RealtimeTTSService:
    """
    实时TTS服务（支持长文本分段合成）
    
    支持将长文本分段，实时合成并流式返回音频
    """
    
    def __init__(self,
                 appid: str = "1871057448",
                 token: str = "tnph6TBEuRv4MWgh425-XUf6K7AcS2_5",
                 cluster: str = "volcano_tts",
                 **kwargs):
        """
        初始化实时TTS服务
        
        Args:
            appid: 应用ID
            token: 访问令牌
            cluster: 集群名称
            **kwargs: 其他参数
        """
        self.tts_service = StreamingTTSService(
            appid=appid,
            token=token,
            cluster=cluster,
            **kwargs
        )
        
        # 分段参数
        self.max_segment_length = kwargs.get("max_segment_length", 500)  # 每段最大字符数
        
        # 用户回调
        self.on_audio_chunk: Optional[Callable] = None  # 音频块回调
        self.on_synthesis_complete: Optional[Callable] = None  # 合成完成回调
        self.on_error: Optional[Callable] = None  # 错误回调
        
        # 设置内部回调
        self.tts_service.on_audio_data = self._on_audio_data
        self.tts_service.on_complete = self._on_complete
        self.tts_service.on_error = self._on_error
        
        print("🎙️ 实时TTS服务初始化完成")
    
    async def _on_audio_data(self, data: Dict[str, Any]):
        """音频数据回调"""
        if self.on_audio_chunk:
            await self.on_audio_chunk(data)
    
    async def _on_complete(self, data: Dict[str, Any]):
        """合成完成回调"""
        if self.on_synthesis_complete:
            await self.on_synthesis_complete(data)
    
    def _on_error(self, error_msg: str):
        """错误回调"""
        if self.on_error:
            self.on_error(error_msg)
    
    def split_text(self, text: str) -> list:
        """
        智能分段文本
        
        Args:
            text: 要分段的文本
        
        Returns:
            文本段列表
        """
        if len(text) <= self.max_segment_length:
            return [text]
        
        segments = []
        current_segment = ""
        
        # 按句子分割（简单实现，按标点符号分割）
        sentences = []
        temp = ""
        for char in text:
            temp += char
            if char in ['。', '！', '？', '.', '!', '?', '\n']:
                sentences.append(temp)
                temp = ""
        if temp:
            sentences.append(temp)
        
        # 组合成合适长度的段落
        for sentence in sentences:
            if len(current_segment) + len(sentence) <= self.max_segment_length:
                current_segment += sentence
            else:
                if current_segment:
                    segments.append(current_segment)
                current_segment = sentence
        
        if current_segment:
            segments.append(current_segment)
        
        return segments
    
    async def synthesize(self, text: str) -> bool:
        """
        合成文本（支持长文本分段）
        
        Args:
            text: 要合成的文本
        
        Returns:
            是否成功
        """
        try:
            # 连接服务
            if not await self.tts_service.connect():
                return False
            
            # 分段处理
            segments = self.split_text(text)
            print(f"📄 文本已分为 {len(segments)} 段")
            
            # 逐段合成
            for i, segment in enumerate(segments):
                print(f"\n🔄 合成第 {i+1}/{len(segments)} 段...")
                success = await self.tts_service.synthesize_text(segment)
                if not success:
                    print(f"❌ 第 {i+1} 段合成失败")
                    return False
                
                # 如果不是最后一段，需要重新连接
                if i < len(segments) - 1:
                    await self.tts_service.disconnect()
                    await asyncio.sleep(0.1)  # 短暂延迟
                    if not await self.tts_service.connect():
                        return False
            
            return True
            
        except Exception as e:
            print(f"❌ 合成异常: {e}")
            if self.on_error:
                self.on_error(str(e))
            return False
        finally:
            await self.tts_service.disconnect()


# 测试函数
async def test_streaming_tts():
    """测试流式TTS"""
    print("=" * 60)
    print("流式TTS测试")
    print("=" * 60)
    
    # 测试文本
    test_text = "你好，欢迎使用火山引擎实时语音合成服务。这是一个测试文本，用来验证TTS服务是否正常工作。"
    
    # 创建服务实例
    tts_service = StreamingTTSService(
        appid="1871057448",
        token="tnph6TBEuRv4MWgh425-XUf6K7AcS2_5",
        cluster="volcano_tts",
        voice_type="BV009_streaming",
        encoding="mp3"
    )
    
    # 保存音频文件
    audio_file = open("test_tts_output.mp3", "wb")
    
    # 音频数据回调
    async def on_audio_data(data):
        audio_file.write(data['audio_data'])
        print(f"📦 收到音频块 {data['chunk_index']}, 大小: {len(data['audio_data'])} 字节")
    
    # 完成回调
    async def on_complete(data):
        audio_file.close()
        print(f"🎉 合成完成！音频已保存到 test_tts_output.mp3")
        print(f"   总块数: {data['total_chunks']}")
        print(f"   总大小: {data['total_size']} 字节")
    
    # 错误回调
    def on_error(error_msg):
        print(f"❌ 错误: {error_msg}")
        audio_file.close()
    
    # 设置回调
    tts_service.on_audio_data = on_audio_data
    tts_service.on_complete = on_complete
    tts_service.on_error = on_error
    
    # 连接并合成
    if await tts_service.connect():
        await tts_service.synthesize_text(test_text)
        await tts_service.disconnect()
    
    print("\n👋 测试结束")


if __name__ == "__main__":
    asyncio.run(test_streaming_tts())

